import re

# 读取 txt 文件中的文本
def read_txt_file(file_path):
    with open(file_path, "r", encoding="utf-8") as f:
        content = f.read()
    return content


def split_text_by_image_tag(text):
    pattern = r"\[image\d+\]"
    segments = re.split(pattern, text)
    segments = [s.strip() for s in segments if s.strip()]
    return segments

def get_subtitles(srt_content):
    pattern = r"(\d+)\n(\d{2}:\d{2}:\d{2},\d{3}) --> (\d{2}:\d{2}:\d{2},\d{3})\n(.+)\n"
    subtitle_matches = re.findall(pattern, srt_content)
    subtitles = []
    for i, match in enumerate(subtitle_matches):
        index = match[0].strip()
        start_time = match[1].strip()
        end_time = match[2].strip()
        subtitle_text = match[3].strip()
        subtitle = [index, start_time, end_time, subtitle_text]
        subtitles.append(subtitle)
    return subtitles


def time_to_ms(time_str):
    time_str = time_str.replace(',', '.')
    hours, minutes, seconds = time_str.split(':')
    seconds, milliseconds = seconds.split('.')
    total_ms = int(hours) * 3600000 + int(minutes) * 60000 + \
        int(seconds) * 1000 + int(milliseconds)
    return total_ms


def ms_to_time_str(ms):
    hours = ms // 3600000
    ms %= 3600000
    minutes = ms // 60000
    ms %= 60000
    seconds = ms // 1000
    ms %= 1000
    return f"{hours:02d}:{minutes:02d}:{seconds:02d},{ms:03d}"


def read_srt_file(file_path):
    with open(file_path, "r", encoding="utf-8") as f:
        content = f.read()
    return content

def time_to_ms(time_str):
    time_str = time_str.replace(',', '.')
    hours, minutes, seconds = time_str.split(':')
    seconds, milliseconds = seconds.split('.')
    total_ms = int(hours) * 3600000 + int(minutes) * 60000 + int(seconds) * 1000 + int(milliseconds)
    return total_ms

def ms_to_time_str(ms):
    hours = ms // 3600000
    ms %= 3600000
    minutes = ms // 60000
    ms %= 60000
    seconds = ms // 1000
    ms %= 1000
    return f"{hours:02d}:{minutes:02d}:{seconds:02d},{ms:03d}"

def milliseconds_to_timestamp(milliseconds):
    total_seconds, mmm = divmod(milliseconds, 1000)
    hh, remainder = divmod(total_seconds, 3600)
    mm, ss = divmod(remainder, 60)
    timestamp = f"{hh:02d}:{mm:02d}:{ss:02d},{mmm:03d}"
    return timestamp

def timestamp_to_milliseconds(timestamp):
    hh, mm, ss_mmm = timestamp.split(':')
    ss, mmm = ss_mmm.split(',')
    milliseconds = int(hh) * 3600000 + int(mm) * \
        60000 + int(ss) * 1000 + int(mmm)
    return milliseconds

# 分段情况
def process_srt(srt_content_list, segments):
    result = []
    str_end = ''  # 上一次剩余的字符
    ti_end = None  # 上一次剩余的时间
    myindex = 0
    for segment in segments:
        myindex += 1
        if myindex == 5:
            print(5)
        srtall = ''  # 本次循环中所有字符
        if str_end:  # 如果上次循环中存在剩余字符，则加到本次字符中
            srtall += str_end
            str_end = ''  # 加上去过后清空
        start_time = srt_content_list[0][1]  # 本次字符开始之前，先记录开始时间
        if ti_end:  # 如果上次还有剩余时长 减掉
            start_time = milliseconds_to_timestamp(
                timestamp_to_milliseconds(start_time)-ti_end)
            ti_end = ''  # 加上去过后清空
        for index, item in enumerate(srt_content_list):
            srtall += item[3].replace(" ", "")
            # 如果包含
            if segment in srtall:
                end_time = ''
                if srtall == segment:  # 包含并且等于的情况
                    end_time = item[2]
                    ti = timestamp_to_milliseconds(
                        end_time)-timestamp_to_milliseconds(start_time)
                    result.append(int(round(ti)))
                    srt_content_list = srt_content_list[index+1:]
                    break
                else:  # 包含超出的情况
                    end_time = item[2]
                    zsc = timestamp_to_milliseconds(
                        end_time)-timestamp_to_milliseconds(start_time)  # 总时长
                    pjsc = zsc/len(srtall)  # 平均每个字符的时长
                    result.append(int(round(pjsc*len(segment))))
                    str_end = srtall.replace(segment, "")  # 补下一段的文本
                    ti_end = int(round(pjsc*len(str_end)))  # 补下一段的时长 下一次开始的时长

                    srt_content_list = srt_content_list[index+1:]
                    break

    return result




txt_file = "srt_txt/txt.txt"
srt_file = "srt_txt/srt.srt"

text = read_txt_file(txt_file)
segments = split_text_by_image_tag(text)

srt_content = read_srt_file(srt_file)
srt_content_list = get_subtitles(srt_content)



result = process_srt(srt_content_list, segments)
print(result)
print(len(result))
with open("srt_txt/片段时长.txt", "w") as file:
    for item in result:
        file.write(str(item) + "\n")
